"""
2.	把inceptionnet-v3预训练模型，作为主干网络，将猫狗数据集图像，映射为2048维向量，然后训练后端全连接网络，进行猫狗图像分类。
按照下述要求，完成相应操作（60分）
【For Tensorflow 1.x】
"""
import tensorflow as tf
import numpy as np
from sklearn.model_selection import train_test_split
import os
import sys

# ①	定义数据参数，包括：图像文件夹、图像特征向量保存地址、inception-v3模型路径和模型参数等
VER = 'v1.1'
IMG_DIR = 'data_zoo'
MODEL_PATH = '../../../../large_data/model/inceptionV3/tensorflow_inception_graph.pb'
INPUT_PLACEHOLDER = 'DecodeJpeg/contents:0'
OUTPUT_TENSOR = 'pool_3/_reshape:0'
ALPHA = 1e-3
N_EPOCHS = 20

np.random.seed(1)
tf.random.set_random_seed(1)
FILE_NAME = os.path.basename(__file__)
LOG_DIR = os.path.join('_log', FILE_NAME, VER)
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
SAVE_DIR_CKPT = os.path.join(SAVE_DIR, 'ckpt')
SAVE_DIR_VEC = os.path.join(SAVE_DIR, 'vectors')

if not os.path.exists(SAVE_DIR_VEC):
    # ②	创建计算图，加载inception-v3预训练模型，返回数据输入张量和瓶颈层输出张量
    graphDef = tf.GraphDef()
    with open(MODEL_PATH, 'rb') as f:
        graphDef.ParseFromString(f.read())
    input_ph, output_tensor = tf.import_graph_def(graphDef, return_elements=[
        INPUT_PLACEHOLDER,
        OUTPUT_TENSOR
    ])

    # ③	开启会话，读取所有的图像，将图像映射的特征向量保存在相应地址
    with tf.Session() as sess:
        with tf.summary.FileWriter(logdir=LOG_DIR) as fw:
            fw.add_graph(sess.graph, global_step=0)
            sess.run(tf.global_variables_initializer())

            os.makedirs(SAVE_DIR_VEC, exist_ok=True)

            cnt = 0
            for sub_dir in os.listdir(IMG_DIR):
                sub_dir_path = os.path.join(IMG_DIR, sub_dir)
                vec_sub_dir_path = os.path.join(SAVE_DIR_VEC, sub_dir)
                os.makedirs(vec_sub_dir_path, exist_ok=True)

                for file in os.listdir(sub_dir_path):
                    file_path = os.path.join(sub_dir_path, file)
                    vec_file_path = os.path.join(vec_sub_dir_path, file)
                    with open(file_path, 'rb') as f:
                        bin = f.read()
                    vec = sess.run(output_tensor, feed_dict={input_ph: bin})
                    vec = np.squeeze(vec)
                    np.savetxt(vec_file_path, vec)
                    cnt += 1
                    if cnt % 25 == 0:
                        print(f'Processed {cnt} pictures ...')

# ④	读取特征向量数据，划分训练集（0.8）、验证集（0.1）、测试集（0.1）
x = []
y = []
yv = 0
for sub_dir in os.listdir(SAVE_DIR_VEC):
    sub_dir_path = os.path.join(SAVE_DIR_VEC, sub_dir)
    for file in os.listdir(sub_dir_path):
        file_path = os.path.join(sub_dir_path, file)
        vec = np.loadtxt(file_path)
        x.append(vec)
        y.append(yv)
    yv += 1
x = np.float32(x)
y = np.int64(y)
y = np.eye(2)[y]
print('x', x.shape)
print('y', y.shape)
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.8, random_state=1, shuffle=True)
x_val, x_test, y_val, y_test = train_test_split(x_test, y_test, train_size=0.5, random_state=1, shuffle=True)
print('x_train', x_train.shape)
print('x_val', x_val.shape)
print('x_test', x_test.shape)
print('y_train', y_train.shape)
print('y_val', y_val.shape)
print('y_test', y_test.shape)

# ⑤	创建全连接神经网络后端模型，实现5种农作物分类任务
x_ph = tf.placeholder(tf.float32, [None, 2048], 'x_ph')
y_ph = tf.placeholder(tf.int64, [None, 2], 'x_ph')
pred = tf.layers.Dense(2, activation=None)(x_ph)

# ⑥	正确合理使用损失函数和优化器
loss = tf.reduce_mean(
    tf.nn.softmax_cross_entropy_with_logits_v2(labels=y_ph, logits=pred)
)
acc = tf.reduce_mean(
    tf.cast(
        tf.equal(
            tf.argmax(y_ph, axis=1),
            tf.argmax(pred, axis=1),
        ),
        tf.float32
    )
)
optim = tf.train.AdamOptimizer(learning_rate=ALPHA).minimize(loss)

# ⑦	训练集数据进行模型训练，自拟合适的超参数
# ⑧	训练1000次，自拟数据批次，每200次打印损失值和交叉验证准确率
group = int(np.ceil(N_EPOCHS / 20))
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    for epoch in range(N_EPOCHS):
        _, lossv, accv = sess.run([optim, loss, acc], feed_dict={x_ph: x_train, y_ph: y_train})

        # ⑨	设置保存检查点，每200次保存后端模型参数
        if epoch % group == 0 or epoch == N_EPOCHS - 1:
            print(f'epoch#{epoch + 1}: loss = {lossv}, acc = {accv}')

    # ⑩	根据测试集，打印测试准确率
    print('Testing...')
    lossv, accv = sess.run([loss, acc], feed_dict={x_ph: x_test, y_ph: y_test})
    print(f'Test loss = {lossv}, acc = {accv}')
    print('Tested')

print('Over')
